﻿// ParseLib - An open source library for natural language processing.
// Copyright (C) Kais Dukes, 2012. E-mail: kais@kaisdukes.com
// http://parselib.org

using System;
using System.Collections.Generic;
using System.Linq;
using log4net;
using ParseLib.Corpora;
using ParseLib.Tagging.Brill.Rules;

namespace ParseLib.Tagging.Brill
{
    public class BrillTrainer
    {
        private static readonly ILog Log = LogManager.GetLogger(typeof (BrillTrainer));
        private BrillTagger _tagger;
        private Sentence[] _sentences;

        public void Train(BrillTagger tagger, IEnumerable<TextualUnit> trainingData)
        {
            // Build the lexicon.
            _tagger = tagger;
            var lexicon = new LexiconBuilder().BuildLexicon(trainingData);
            _tagger.Lexicon = lexicon;

            // Initial tagging.
            _sentences = trainingData.Select(x => new Sentence(x, lexicon)).ToArray();

            // Learn rules.
            LearnRules();
        }

        private void LearnRules()
        {
            Tuple<Rule, int> pair;
            while (IsValidRule(pair = GetBestRule()))
            {
                // Add the rule.
                var rule = pair.Item1;
                var score = pair.Item2;
                Log.DebugFormat("Adding rule: {0} (Score = {1})", rule, score);
                _tagger.Rules.Add(rule);

                // Apply the rule.
                ApplyRule(rule);
            }
        }

        private void ApplyRule(Rule rule)
        {
            Log.Debug("Applying rule...");
            foreach (var sentence in _sentences)
            {
                rule.Apply(sentence);
            }
        }

        private Tuple<Rule, int> GetBestRule()
        {
            // Rule types.
            var types = new[]
                            {
                                RuleType.Previous,
                                //RuleType.OneOfPreviousTwo,
                                //RuleType.OneOfPreviousThree
                            };

            // Find best rule.
            Rule bestRule = null;
            var bestScore = 0;
            foreach (var type in types)
            {
                var pair = GetBestRule(type);
                var rule = pair.Item1;
                var score = pair.Item2;
                if (score <= bestScore) continue;
                bestRule = rule;
                bestScore = score;
            }

            // Return the best rule (with score).
            return bestRule != null ? new Tuple<Rule, int>(bestRule, bestScore) : null;
        }

        private Tuple<Rule, int> GetBestRule(RuleType type)
        {
            // Find rule parameters.
            Log.DebugFormat("Finding rule parameters for {0}...", Rule.FormatRuleType(type));
            var rules = new HashSet<Rule>();
            foreach (var sentence in _sentences)
            {
                var size = sentence.TokenCount;
                for (var i = 0; i < size; i++)
                {
                    // Correct tag?
                    var token = sentence[i];
                    if (token.IsCorrect) continue;

                    // Context.
                    AddRulesFromContext(rules, type, sentence, i);
                }
            }

            // Search for best rule.
            Log.DebugFormat("Searching for best rule out of {0}...", rules.Count);
            Rule bestRule = null;
            var bestScore = 0;
            foreach (var rule in rules)
            {
                var score = EvaluateRule(rule);
                if (score <= bestScore) continue;
                bestRule = rule;
                bestScore = score;
            }

            // Return the best rule (with score).
            return bestRule != null ? new Tuple<Rule, int>(bestRule, bestScore) : null;
        }

        private static void AddRulesFromContext(
            ISet<Rule> rules, RuleType type, Sentence sentence, int index)
        {
            var t0 = sentence[index];
            var p1 = sentence[index - 1];
            var p2 = sentence[index - 2];
            var p3 = sentence[index - 3];

            switch (type)
            {
                case RuleType.Previous:
                    rules.Add(GetRule(type, t0, p1));
                    break;
                case RuleType.OneOfPreviousTwo:
                    rules.Add(GetRule(type, t0, p1));
                    rules.Add(GetRule(type, t0, p2));
                    break;
                case RuleType.OneOfPreviousThree:
                    rules.Add(GetRule(type, t0, p1));
                    rules.Add(GetRule(type, t0, p2));
                    rules.Add(GetRule(type, t0, p3));
                    break;
                default:
                    throw new NotImplementedException();
            }
        }

        private static SingleTagRule GetRule(RuleType type, BrillToken token, BrillToken context)
        {
            return new SingleTagRule
                       {
                           RuleType = type,
                           From = token.CurrentTag,
                           To = token.CorrectTag,
                           Context = context != null ? context.CurrentTag : null
                       };
        }

        private int EvaluateRule(Rule rule)
        {
            var score = 0;
            foreach (var sentence in _sentences)
            {
                var size = sentence.TokenCount;
                for (var i = 0; i < size; i++)
                {
                    if (rule.Match(sentence, i))
                    {
                        score += rule.To == sentence[i].CorrectTag ? 1 : -1;
                    }
                }
            }
            return score;
        }

        private bool IsValidRule(Tuple<Rule, int> pair)
        {
            return pair != null && pair.Item2 >= _tagger.Threshold;
        }
    }
}